package com.data.center.util;

import java.io.FileOutputStream;
import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import org.apache.poi.ss.usermodel.*;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;

public class DatabaseToExcelUtils {

    public static void main(String[] args) {
        String url = "jdbc:mysql://localhost:3306/data_center";
        String user = "root";
        String password = "123456";
        String tableName = "person"; //需要抽取数据的数据库中表的名称
        int numRowsToRetrieve = 200; // 从数据库中抽取的行数

        try (Connection conn = DriverManager.getConnection(url, user, password)) { // 建立数据库连接
            List<String[]> data = retrieveRandomRows(conn, tableName, numRowsToRetrieve); // 从数据库中随机抽取数据
            writeToExcel(data, ".\\source\\Person\\Person测试用例随机抽取.xlsx"); // 将数据写入Excel文件，使用相对路径保存
            System.out.println("Data exported to Excel successfully."); // 输出成功信息
        } catch (SQLException | IOException e) { // 捕获SQL和IO异常
            e.printStackTrace(); // 打印异常信息
        }
    }

    private static List<String[]> retrieveRandomRows(Connection conn, String tableName, int numRowsToRetrieve) throws SQLException {
        List<String[]> data = new ArrayList<>(); // 创建一个存储数据的列表
        String sql = "SELECT * FROM " + tableName + " ORDER BY RAND() LIMIT ?"; // SQL查询语句，随机抽取指定数量的行
        try (PreparedStatement stmt = conn.prepareStatement(sql)) { // 创建预编译的SQL语句
            stmt.setInt(1, numRowsToRetrieve); // 设置查询结果行数参数
            ResultSet rs = stmt.executeQuery(); // 执行查询
            ResultSetMetaData metaData = rs.getMetaData(); // 获取结果集的元数据
            int columnCount = metaData.getColumnCount(); // 获取结果集的列数
            String[] columnNames = new String[columnCount]; // 创建一个数组来存储列名
            for (int i = 1; i <= columnCount; i++) { // 从1开始，因为列索引从1开始
                columnNames[i - 1] = metaData.getColumnName(i); // 获取列名并存储到数组中
            }
            data.add(columnNames); // 将列名添加到数据列表中
            while (rs.next()) { // 遍历查询结果集
                // 读取每一行数据并添加到列表中
                String[] row = new String[columnCount];
                for (int i = 0; i < columnCount; i++) {
                    row[i] = rs.getString(i + 1);
                }
                data.add(row); // 将每一行数据添加到列表中
            }
        }
        return data; // 返回抽取的数据列表
    }

    private static void writeToExcel(List<String[]> data, String filePath) throws IOException {
        try (Workbook workbook = new XSSFWorkbook()) { // 创建Excel工作簿对象
            Sheet sheet = workbook.createSheet("Sheet1"); // 创建Excel表格页
            int rowNum = 0; // 行数计数器
            for (String[] rowData : data) { // 遍历数据列表中的每一行数据
                Row row = sheet.createRow(rowNum++); // 创建Excel行对象
                int colNum = 0; // 列数计数器
                for (String cellData : rowData) { // 遍历每一行数据中的每个单元格
                    row.createCell(colNum++).setCellValue(cellData); // 创建单元格并设置单元格值
                }
            }
            // 将Excel文件写入磁盘
            try (FileOutputStream outputStream = new FileOutputStream(filePath)) {
                workbook.write(outputStream); // 写入Excel文件
            }
        }
    }
}
